#pragma once
#include <iostream>
#include <utility>
using namespace std;

enum COLOR
{
    RED,
    BLACK,
};

template<class V>
struct __rb_tree_node
{
    __rb_tree_node<V>* _left;
    __rb_tree_node<V>* _right;
    __rb_tree_node<V>* _parent;

    V _val;
    COLOR _col;

    __rb_tree_node<V>(const V& val)
        : _left(nullptr), _right(nullptr), _parent(nullptr)
        , _val(val), _col(RED)
    {}
};

template<class V, class Ref, class Ptr>
struct __rb_tree_iterator
{
    typedef __rb_tree_node<V> node;
    typedef __rb_tree_iterator<V, V&, V*> iterator;
    typedef __rb_tree_iterator<V, Ref, Ptr> self;
    node* _node = nullptr;

    __rb_tree_iterator<V, Ref, Ptr>(node* node)
        : _node(node)
    {}

    // support normal iter construt const iter
    __rb_tree_iterator(const __rb_tree_iterator<V, V&, V*>& it)
        : _node(it._node)
    {}

    Ref operator*()
    {
        return _node->_val;
    }

    Ptr operator->()
    {
        return &_node->_val;
    }

    bool operator==(const self& s)
    {
        return _node == s._node;
    }

    bool operator!=(const self& s)
    {
        return _node != s._node;
    }

    self& operator++()
    {
        if (_node->_right)
        {
            node* left = _node->_right;
            while (left->_left)
            {
                left = left->_left;
            }
            _node = left;
        }
        else
        {
            node* cur = _node;
            node* parent = cur->_parent;

            while (parent && parent->_right == cur)
            {
                cur = parent;
                parent = cur->_parent;
            }
            _node = parent;
        }
        return *this;
    }

    self& operator--()
    {
        if (!_node) assert(false); // can't do that

        if (_node->_left)
        {
            node* right = _node->_left;
            while (right->_right)
            {
                right = right->_right;
            }
            _node = right;
        }
        else
        {
            node* cur = _node;
            node* parent = cur->_parent;

            while (parent && parent->_left == cur)
            {
                cur = parent;
                parent = cur->_parent;
            }
            _node = parent;
        }
        return *this;
    }


};

template<class K, class V, class KeyOfValue, class Compare>
class rb_tree
{
public:
    typedef __rb_tree_node<V> node;
    typedef __rb_tree_iterator<V, V&, V*> iterator;
    typedef __rb_tree_iterator<V, const V&, const V*> const_iterator;

    iterator begin()
    {
        node* cur = _root;
        while (cur && cur->_left) {
            cur = cur->_left;
        }
        return iterator(cur);
    }

    iterator end()
    {
        return iterator(nullptr);
    }

public:
    rb_tree()
    {}

    rb_tree(const rb_tree& t)
    {
        _root = copy(_root);
    }

    rb_tree& operator=(const rb_tree t)
    {
        if (this != &t)
        {
            swap(_root, t._root);
        }
        return *this;
    }

    ~rb_tree()
    {
        destroy(_root);
        _root = nullptr;
    }

private:
    node* copy(node* root)
    {
        if (!root)
            return nullptr;

        node* new_node = new node(root->_kv);
        new_node->_left = copy(root->_left);
        new_node->_right = copy(root->_right);
        return new_node;
    }

    void destroy(node* root)
    {
        if (!root)
            return;

        destroy(root->_left);
        destroy(root->_right);
        delete root;
    }

public:
    iterator find(const K& key)
    {
        node* cur = _root;

        while (cur)
        {
            if (_cmp(_kov(cur->_val), key))
                cur = cur->_right;
            else if (_cmp(_kov(cur->_val), key))
                cur = cur->_left;
            else
                return iterator(cur);
        }

        return iterator(nullptr);
    }

    pair<iterator, bool> insert(const V& val)
    {
        if (!_root)
        {
            _root = new node(val);
            _root->_col = BLACK;
            return {iterator(_root), true};
        }

        node* parent = nullptr;
        node* cur = _root;

        while (cur)
        {
            if (_cmp(_kov(cur->_val), _kov(val)))
            {
                parent = cur;
                cur = cur->_right;
            }
            else if (_cmp(_kov(val), _kov(cur->_val)))
            {
                parent = cur;
                cur = cur->_left;
            }
            else
            {
                return {iterator(cur), false};
            }
        }

        cur = new node(val);
        if (_cmp(_kov(parent->_val), _kov(val)))
            parent->_right = cur;
        else
            parent->_left = cur;
        cur->_parent = parent;

        while (parent && parent->_col == RED)
        {
            node* grandpa = parent->_parent;

            if (grandpa->_left == parent)
            {
                node* uncle = grandpa->_right;

                if (uncle && uncle->_col == RED)
                {
                    parent->_col = uncle->_col = BLACK;
                    grandpa->_col = RED;

                    cur = grandpa;
                    parent = cur->_parent;
                }
                else
                {
                    if (parent->_left == cur)
                    {
                        rotate_r(grandpa);
                        cur->_col = grandpa->_col = RED;
                        parent->_col = BLACK;
                    }
                    else
                    {
                        rotate_l(parent);
                        rotate_r(grandpa);
                        grandpa->_col = parent->_col = RED;
                        cur->_col = BLACK;
                    }
                    break;
                }
            }
            else
            {
                node* uncle = grandpa->_left;

                if (uncle && uncle->_col == RED)
                {
                    parent->_col = uncle->_col = BLACK;
                    grandpa->_col = RED;

                    cur = grandpa;
                    parent = cur->_parent;
                }
                else
                {
                    if (parent->_left == cur)
                    {
                        rotate_r(parent);
                        rotate_l(grandpa);
                        grandpa->_col = parent->_col = RED;
                        cur->_col = BLACK;
                    }
                    else
                    {
                        rotate_l(grandpa);
                        cur->_col = grandpa->_col = RED;
                        parent->_col = BLACK;
                    }
                    break;
                }
            }
        }

        _root->_col = BLACK;
        return {iterator(cur), true};
    }

private:
    void rotate_l(node* parent)
    {
        node* subr = parent->_right;
        node* subrl = subr->_left;

        parent->_right = subrl;
        if (subrl) subrl->_parent = parent;

        node* pparent = parent->_parent;

        subr->_left = parent;
        parent->_parent = subr;

        if (parent == _root)
            _root = subr;
        else
        {
            if (pparent->_left == parent)
                pparent->_left = subr;
            else
                pparent->_right = subr;
        }
        subr->_parent = pparent;
    }

    void rotate_r(node* parent)
    {
        node* subl = parent->_left;
        node* sublr = subl->_right;

        parent->_left = sublr;
        if (sublr) sublr->_parent = parent;

        node* pparent = parent->_parent;

        subl->_right = parent;
        parent->_parent = subl;

        if (parent == _root)
        {
            _root = subl;
            subl->_parent = nullptr;
        }
        else
        {
            if (pparent->_left == parent)
                pparent->_left = subl;
            else
                pparent->_right = subl;
            subl->_parent = pparent;
        }
    }

public:
    void inorder()
    {
        _inorder(_root);
        cout << endl;
    }

    bool is_rbtree()
    {
        if (_root && _root->_col == RED)
            return false;

        int mark = -1;
        return check(_root, mark, 0);
    }

private:
    void _inorder(node* root)
    {
        if (!root) return;

        _inorder(root->_left);
        cout << _kov(root->_val) << ":" << _kov(root->_val) << " ";
        _inorder(root->_right);
    }

    bool check(node* root, int& mark, int cnt)
    {
        if (!root)
        {
            if (mark == -1)
                mark = cnt;
            else if (mark != cnt)
            {
                cout << "block nodes count error\n" << endl;
                return false;
            }
            return true;
        }

        if (root->_col == BLACK)
            cnt++;

        if (root->_col == RED && root->_parent->_col == RED)
        {
            cout << "consecutive red nodes "
                 << _kov(root->_parent->_val) << " and "
                 << _kov(root->_val) << endl;
            return false;
        }

        return check(root->_left, mark, cnt)
               && check(root->_right, mark, cnt);
    }

private:
    KeyOfValue _kov;
    Compare _cmp;

    node* _root = nullptr;
};
